# encoding: utf-8
"""
@author: Yuxian Meng
@contact: yuxian_meng@shannonai.com

@version: 1.0
@file: run_rcnn
@time: 2020/12/31 14:34
@desc: 

"""


# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# Requires vqa-maskrcnn-benchmark to be built and installed. See Readme
# for more details.
import argparse
import glob
import os

import cv2
import numpy as np
import torch
from PIL import Image
from tqdm import tqdm
from typing import List
import logging
import os
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.layers import nms
from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.structures.image_list import to_image_list
from maskrcnn_benchmark.utils.model_serialization import load_state_dict
from rcnn.dataset import get_dataloader


class FeatureExtractor:
    MAX_SIZE = 1333
    MIN_SIZE = 800

    def __init__(self, args):
        self.args = args
        os.environ["CUDA_VISIBLE_DEVICES"] = str(self.args.cuda)
        self.detection_model = self._build_detection_model()

        os.makedirs(self.args.output_folder, exist_ok=True)

    @staticmethod
    def get_parser():
        parser = argparse.ArgumentParser()
        parser.add_argument(
            "--model_file", default=None, type=str, help="Detectron model file"
        )
        parser.add_argument(
            "--config_file", default=None, type=str, help="Detectron config file"
        )
        parser.add_argument("--batch_size", type=int, default=1, help="Batch size")
        parser.add_argument(
            "--num_features",
            type=int,
            default=100,
            help="Number of features to extract.",
        )
        parser.add_argument(
            "--output_folder", type=str, default="./output", help="Output folder"
        )
        parser.add_argument("--image_dir", type=str, help="Image directory or file")
        parser.add_argument(
            "--feature_name",
            type=str,
            help="The name of the feature to extract",
            default="fc6",
        )
        parser.add_argument(
            "--confidence_threshold",
            type=float,
            default=0,
            help="Threshold of detection confidence above which boxes will be selected",
        )
        parser.add_argument(
            "--background",
            action="store_true",
            help="The model will output predictions for the background class when set",
        )
        parser.add_argument(
            "--partition", type=int, default=0, help="Partition to download."
        )
        parser.add_argument(
            "--workers", type=int, default=1, help="cpu workers"
        )
        parser.add_argument(
            "--start", type=int, default=-1, help="start offset"
        )
        parser.add_argument(
            "--end", type=int, default=-1, help="end offset"
        )
        parser.add_argument(
            "--num_gpus", type=int, default=8, help="num_gpus"
        )
        parser.add_argument(
            "--workers_per_gpu", type=int, default=4, help="worker per gpu"
        )
        parser.add_argument(
            "--cuda", type=int, default=0, help="cuda_device"
        )
        return parser

    def _build_detection_model(self):
        cfg.merge_from_file(self.args.config_file)
        cfg.freeze()

        model = build_detection_model(cfg)
        checkpoint = torch.load(self.args.model_file, map_location=torch.device("cpu"))

        load_state_dict(model, checkpoint.pop("model"))

        model.to("cuda")
        model.eval()
        return model

    def _image_transform(self, path):
        img = Image.open(path)
        im = np.array(img).astype(np.float32)
        # IndexError: too many indices for array, grayscale images
        if len(im.shape) < 3:
            im = np.repeat(im[:, :, np.newaxis], 3, axis=2)
        im = im[:, :, ::-1]
        im -= np.array([102.9801, 115.9465, 122.7717])
        im_shape = im.shape
        im_height = im_shape[0]
        im_width = im_shape[1]
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])

        # Scale based on minimum size
        im_scale = self.MIN_SIZE / im_size_min

        # Prevent the biggest axis from being more than max_size
        # If bigger, scale it down
        if np.round(im_scale * im_size_max) > self.MAX_SIZE:
            im_scale = self.MAX_SIZE / im_size_max

        im = cv2.resize(
            im, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR
        )
        img = torch.from_numpy(im).permute(2, 0, 1)

        im_info = {"width": im_width, "height": im_height}

        return img, im_scale, im_info

    def _process_feature_extraction(
        self, output, im_scales, im_infos, feature_name="fc6", conf_thresh=0
    ):
        batch_size = len(output[0]["proposals"])
        n_boxes_per_image = [len(boxes) for boxes in output[0]["proposals"]]
        score_list = output[0]["scores"].split(n_boxes_per_image)
        score_list = [torch.nn.functional.softmax(x, -1) for x in score_list]
        feats = output[0][feature_name].split(n_boxes_per_image)
        cur_device = score_list[0].device

        feat_list = []
        info_list = []

        for i in range(batch_size):
            dets = output[0]["proposals"][i].bbox / im_scales[i]
            scores = score_list[i]
            max_conf = torch.zeros((scores.shape[0])).to(cur_device)
            conf_thresh_tensor = torch.full_like(max_conf, conf_thresh)
            start_index = 1
            # Column 0 of the scores matrix is for the background class
            if self.args.background:
                start_index = 0
            for cls_ind in range(start_index, scores.shape[1]):
                cls_scores = scores[:, cls_ind]
                keep = nms(dets, cls_scores, 0.5)
                max_conf[keep] = torch.where(
                    # Better than max one till now and minimally greater than conf_thresh
                    (cls_scores[keep] > max_conf[keep])
                    & (cls_scores[keep] > conf_thresh_tensor[keep]),
                    cls_scores[keep],
                    max_conf[keep],
                )

            sorted_scores, sorted_indices = torch.sort(max_conf, descending=True)
            num_boxes = (sorted_scores[: self.args.num_features] != 0).sum()
            keep_boxes = sorted_indices[: self.args.num_features]
            feat_list.append(feats[i][keep_boxes])
            bbox = output[0]["proposals"][i][keep_boxes].bbox / im_scales[i]
            # Predict the class label using the scores
            objects = torch.argmax(scores[keep_boxes][start_index:], dim=1)
            cls_prob = torch.max(scores[keep_boxes][start_index:], dim=1)

            info_list.append(
                {
                    "bbox": bbox.cpu().numpy(),
                    "num_boxes": num_boxes.item(),
                    "objects": objects.cpu().numpy(),
                    "image_width": im_infos[i]["width"],
                    "image_height": im_infos[i]["height"],
                    "cls_prob": scores[keep_boxes].cpu().numpy(),
                }
            )

        return feat_list, info_list

    def get_detectron_features(self, image_paths):
        img_tensor, im_scales, im_infos = [], [], []

        for image_path in image_paths:
            im, im_scale, im_info = self._image_transform(image_path)
            img_tensor.append(im)
            im_scales.append(im_scale)
            im_infos.append(im_info)

        # Image dimensions should be divisible by 32, to allow convolutions
        # in detector to work
        current_img_list = to_image_list(img_tensor, size_divisible=32)
        current_img_list = current_img_list.to("cuda")

        with torch.no_grad():
            output = self.detection_model(current_img_list)

        feat_list = self._process_feature_extraction(
            output,
            im_scales,
            im_infos,
            self.args.feature_name,
            self.args.confidence_threshold,
        )

        return feat_list

    def _chunks(self, array, chunk_size):
        for i in range(0, len(array), chunk_size):
            yield array[i : i + chunk_size]

    def _save_feature(self, file_name, feature, info):
        file_base_name = os.path.basename(file_name)
        file_base_name = file_base_name.split(".")[0]
        info["image_id"] = file_base_name
        info["features"] = feature.cpu().numpy()
        output_path = file_name.replace(self.args.image_dir,
                                        self.args.output_folder)
        # file_base_name = file_base_name + ".npy"
        # output_path = os.path.join(self.args.output_folder, file_base_name)
        np.save(output_path, info)
        # np.save(output_path, feature.cpu().numpy())

    def extract_features(self, files, verbose=False):
        # image_dir = self.args.image_dir
        # if os.path.isfile(image_dir):
        #     features, infos = self.get_detectron_features([image_dir])
        #     self._save_feature(image_dir, features[0], infos[0])
        # else:
        """
        files = glob.glob(os.path.join(image_dir, "*.jpg"))
        files = sorted(files)
        # files = [files[i: i+1000] for i in range(0, len(files), 1000)][self.args.partition]
        # 改造为dataset
        # for chunk in tqdm(self._chunks(files, self.args.batch_size)):
        #     try:
        #         features, infos = self.get_detectron_features(chunk)
        #         for idx, file_name in enumerate(chunk):
        #             self._save_feature(file_name, features[idx], infos[idx])
        #     except BaseException as e:
        #         print(e)
        #         continue

        data_loader = get_dataloader(files, batch_size=self.args.batch_size, workers=self.args.workers)
        for chunk in tqdm(data_loader):
            try:
                chunk, im_scales, im_infos, chunk_names = chunk
                chunk = chunk.to("cuda")
                with torch.no_grad():
                    output = self.detection_model(chunk)

                features, infos = self._process_feature_extraction(
                    output,
                    im_scales,
                    im_infos,
                    self.args.feature_name,
                    self.args.confidence_threshold,
                )
                for idx, file_name in enumerate(chunk_names):
                    self._save_feature(file_name, features[idx], infos[idx])
            except BaseException as e:
                logging.error(e, exc_info=True)
                continue
        """
        # files = glob.glob(os.path.join(image_dir, "*/*.jpg"))
        # files = [files[i: i+1000] for i in range(0, len(files), 1000)][self.args.partition]
        # 改造为dataset
        # for chunk in tqdm(self._chunks(files, self.args.batch_size)):
        #     try:
        #         features, infos = self.get_detectron_features(chunk)
        #         for idx, file_name in enumerate(chunk):
        #             self._save_feature(file_name, features[idx], infos[idx])
        #     except BaseException as e:
        #         print(e)
        #         continue

        data_loader = get_dataloader(files, batch_size=self.args.batch_size, workers=self.args.workers)
        iterator = tqdm(data_loader) if verbose else data_loader
        for chunk in iterator:
            try:
                chunk, im_scales, im_infos, chunk_names = chunk
                chunk = chunk.to("cuda")
                with torch.no_grad():
                    output = self.detection_model(chunk)

                features, infos = self._process_feature_extraction(
                    output,
                    im_scales,
                    im_infos,
                    self.args.feature_name,
                    self.args.confidence_threshold,
                )
                for idx, file_name in enumerate(chunk_names):
                    self._save_feature(file_name, features[idx], infos[idx])
            except BaseException as e:
                logging.error(e, exc_info=True)
                continue


def extract_features(args, files: List[str], verbose=False):
    extractor = FeatureExtractor(args)
    extractor.extract_features(files, verbose=verbose)


if __name__ == "__main__":
    from multiprocessing import Process
    args = FeatureExtractor.get_parser().parse_args()
    files = glob.glob(os.path.join(args.image_dir, "*/*.jpg"))
    num_gpus = args.num_gpus
    worker_per_gpu = args.workers_per_gpu
    num_procs = num_gpus * worker_per_gpu

    start = 0 if args.start == -1 else max(args.start, 0)
    end = len(files) if args.end == -1 else min(len(files), args.end)
    files = files[start: end]
    files = [f for f in files if not os.path.exists(f+".npy")]
    chunk_size = (end - start) // num_procs
    print(f"Preprocessing {len(files)} with {num_procs} processes")
    procs = []
    offset = 0
    for cuda_device in range(num_gpus):
        for worker_idx in range(worker_per_gpu):
            proc_end = offset + chunk_size
            verbose = False
            if cuda_device == num_gpus-1 and worker_idx == worker_per_gpu-1:
                proc_end = end
                verbose = True
            proc_files = files[offset: proc_end]
            args.cuda = cuda_device
            p = Process(target=extract_features, kwargs={"args": args, "files": proc_files, "verbose": verbose})
            offset += chunk_size
            p.start()
            procs.append(p)
    for p in procs:
        p.join()
